from langchain_community.vectorstores import Chroma, FAISS
from langchain_core.example_selectors import SemanticSimilarityExampleSelector
from langchain_ollama import OllamaEmbeddings
from fewshot.few_shot_demo import examples

"""
 pip install langchain_Chroma
"""

# 2. 初始化嵌入模型 - 使用本地的 nomic-embed-text
embeddings = OllamaEmbeddings(
    model="nomic-embed-text",  # nomic-embed-text: 专门为检索优化，质量比通用聊天模型好
    base_url="http://localhost:11434"  # Ollama服务的地址
)

example_selector = SemanticSimilarityExampleSelector.from_examples(
    examples,   # 可供选择的示例列表
    # 用于生成嵌入的嵌入类，用户测量 “语义相似性”。文本图片音频转换为向量。 相似度匹配：余弦相似度s
    # OpenAIEmbeddings(),  # 需要设置OPENAI_API_KEY环境变量,默认使用的是"text-embedding-ada-002"嵌入模型
    embeddings=embeddings,
    # 用于存储嵌入并进行相似度搜索的VectorStore类
    vectorstore_cls=Chroma,
    k=2 # 选择最相似的2个示例
)
question = "玛丽·波尔·华盛顿的父亲是谁？"
selected_examples = example_selector.select_examples({"question": question})
print(f"与输入相似的示例：{question}")
for example in selected_examples:
    print("\n")
    for k, v in example.items():
        print(f"{k}: {v}")
